# TODO: 导入必要的库和模块
import gradio as gr
import pickle
import numpy as np
from PIL import Image
import sklearn
from sklearn.datasets import load_digits

print(f"Gradio version: {gr.__version__}")
print(f"Scikit-learn version: {sklearn.__version__}")

# TODO: 加载保存的最佳KNN模型
try:
    with open('best_knn_model.pkl', 'rb') as f:
        best_knn_model = pickle.load(f)
    print("模型加载成功！")
except FileNotFoundError:
    print("错误：找不到 best_knn_model.pkl 文件")
    print("请先运行 optimal_knn.py 来训练和保存模型")
    exit(1)

# 验证模型有效性
digits = load_digits()
print(f"训练数据基本信息：{digits.images[0].shape}像素，像素值范围{digits.data.min():.1f}~{digits.data.max():.1f}")

# 测试模型对前几个训练样本的预测
print("模型对训练样本的预测结果（验证模型有效性）：")
for i in range(5):
    sample = digits.data[i].reshape(1, -1)
    prediction = best_knn_model.predict(sample)[0]
    true_label = digits.target[i]
    print(f"真实数字：{true_label}，模型预测：{prediction}")

def preprocess_image(image):
    """
    预处理手写数字图像，使其与MNIST数据集格式兼容
    """
    try:
        # 将图像转换为PIL Image对象
        if isinstance(image, str):
            # 如果是文件路径
            image = Image.open(image)
        elif isinstance(image, np.ndarray):
            # 如果是numpy数组
            image = Image.fromarray((image * 255).astype('uint8'))
        
        # 将图像转换为灰度
        if image.mode != 'L':
            image = image.convert('L')
        
        # 调整大小为8x8像素（与MNIST数据集相同）
        image = image.resize((8, 8), Image.LANCZOS)
        
        # 转换为numpy数组
        img_array = np.array(image)
        
        # 反转颜色（MNIST是白字黑底，而手写板通常是黑字白底）
        img_array = 255 - img_array
        
        # 归一化到0-16范围（与原始MNIST数据集一致）
        img_array = img_array / 255.0 * 16
        img_array = np.clip(img_array, 0, 16)
        
        # 展平为1D数组
        img_flat = img_array.flatten()
        
        return img_flat.reshape(1, -1)
    
    except Exception as e:
        print(f"图像预处理错误: {e}")
        return None

def predict_digit(image):
    """
    使用KNN模型预测手写数字
    """
    if image is None:
        return "请绘制一个数字"
    
    try:
        # 预处理图像
        processed_image = preprocess_image(image)
        
        if processed_image is None:
            return "图像处理失败，请重试"
        
        # 使用模型进行预测
        prediction = best_knn_model.predict(processed_image)
        probability = best_knn_model.predict_proba(processed_image)
        
        # 获取置信度
        confidence = np.max(probability)
        
        result = f"预测数字: {prediction[0]}\n置信度: {confidence:.2%}"
        
        # 显示所有数字的概率
        result += "\n\n所有数字概率:"
        for i, prob in enumerate(probability[0]):
            result += f"\n数字 {i}: {prob:.2%}"
        
        return result
    
    except Exception as e:
        return f"预测错误: {str(e)}"

# 创建Gradio界面 - 适用于Gradio 3.x版本
def create_webapp():
    # 使用简单的Image输入组件
    interface = gr.Interface(
        fn=predict_digit,
        inputs=gr.inputs.Image(
            image_mode="L",  # 灰度图像
            source="canvas",  # 画布模式
            shape=(200, 200),  # 画布大小
            invert_colors=True,  # 反转颜色
            tool="pencil",  # 铅笔工具
            label="在手写板上绘制数字（0-9）"
        ),
        outputs=gr.outputs.Textbox(
            label="预测结果",
            type="auto"
        ),
        title="手写数字识别 - KNN分类器",
        description="使用训练好的KNN模型识别手写数字。请在下方手写板上绘制一个数字（0-9），模型会自动识别。",
        live=True,  # 实时预测
        examples=[
            ["example_0.png"],  # 你可以添加示例图片
            ["example_1.png"],
        ] if False else None  # 暂时禁用示例
    )
    
    return interface

if __name__ == "__main__":
    try:
        # 创建并启动web应用
        print("创建Web应用...")
        app = create_webapp()
        print("启动Web应用...")
        print("请在浏览器中打开显示的本地地址（通常是 http://127.0.0.1:7860）")
        app.launch(share=False, server_port=7860)
    except Exception as e:
        print(f"创建Web应用失败: {e}")
        print("尝试备用方案...")
        
        # 备用方案：使用更简单的界面
        try:
            simple_interface = gr.Interface(
                fn=predict_digit,
                inputs="sketch",
                outputs="text",
                title="手写数字识别",
                description="绘制0-9的数字进行识别"
            )
            simple_interface.launch(share=False)
        except Exception as e2:
            print(f"备用方案也失败: {e2}")
            print("请检查依赖包版本")